from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
import mydataPF as mydatrg
import Functions as func

def i_inj(t):
    if t <= 1.1*ms and t > 1 *ms:
        return 0.8 *nA
#    elif t <= 60.1*ms and t > 60 *ms:
#        return 0.3 *nA
#    elif t <= 120.1*ms and t > 120 *ms:
#        return 0.2 *nA
    else:
        return 0*nA
    
def current_inj(t):
    if t > 0*msecond:
        return -0.5*namp
    else:
        return 0*mamp
    
class NeuronFactory:
    
    def generate(self, size, init='def'):
        md = mydatrg.PFVars()
        
        # The model
        eqsPF=Equations('''
        dvm/dt = (i_inj(t)-(md.gna*mna*mna*mna*hna*(vm-md.Ena))-(md.gk*mk*mk*mk*mk*(vm-md.Ek))-(md.gnap*mnap*hnap*(vm-md.Ena))-md.gl*(vm-md.El)+gi_current+ge_current)/md.cpf : mvolt
        dmna/dt = (func.m_inf_na(vm/mvolt)-mna)/md.taumna : 1
        dhna/dt = (func.h_inf_na(vm/mvolt)-hna)/func.tau_h_na(vm/mvolt) : 1
        dmk/dt = (func.m_inf_k(vm/mvolt)-mk)/func.tau_m_k(vm/mvolt) : 1
        dmnap/dt = (func.m_inf_nap(vm/mvolt)-mnap)/md.tau_m_nap : 1
        dhnap/dt = (func.h_inf_nap(vm/mvolt)-hnap)/func.tau_h_nap(vm/mvolt) : 1
        dge/dt = -ge*md.invtau : siemens
        dgi/dt = -gi*md.invtau : siemens
        dipsp/dt = (gi-ipsp)*md.invtau : siemens
        depsp/dt = (ge-epsp)*md.invtau : siemens
        gi_current = ipsp*(md.Ei-vm) : amp
        ge_current = epsp*(md.Ee-vm) : amp
        ''')
        
        
        P=NeuronGroup(size,model=eqsPF,
        threshold=EmpiricalThreshold(threshold=-40*mV),
        implicit=True)
        # Initialization
        
        vv =md.vinit
        if init != 'def':
            vv=init
        P.vm=vv*mV 
        P.mna = func.m_inf_na(vv)
        P.hna = func.h_inf_na(vv)
        P.mk = func.m_inf_k(vv)
        
        P.mnap = func.m_inf_nap(vv)
        P.hnap = func.h_inf_nap(vv)
        P.El = -64*mV * randn(len(P))*0.64*mV
        return P